620E - New Year Tree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 4e5+10;
typedef long long ll;
typedef pair<ll, ll> pii;
#define pb push_back
#define sum first
#define lz second

int t, n, q, v[MAXN];
int euler[MAXN], tin[MAXN], tout[MAXN];
pii seg[4*MAXN];
vector<int> edge[MAXN];

void dfs(int x, int pai) {
   t++;
   euler[t] = v[x];
   tin[x] = t;
   for(int viz : edge[x])
      if(viz!=pai) dfs(viz, x);
   tout[x] = t;
}

void build(int node, int l, int r) {
   if(l==r) {
      seg[node].sum = ((ll)1<<euler[l]);
      return;
   }
   int meio = (l+r)/2;
   build(2*node, l, meio);
   build(2*node+1, meio+1, r);
   seg[node].sum = seg[2*node].sum | seg[2*node+1].sum;
}

void lazy(int node, int l, int r) {
   if(seg[node].lz==0) return;
   seg[node].sum = ((ll)1<<seg[node].lz);
   if(l!=r)
      seg[2*node].lz = seg[2*node+1].lz = seg[node].lz;
   seg[node].lz = 0;
}

void update(int node, int l, int r, int p, int q, int x) {
   lazy(node, l, r);
   if(r<p or q<l) return;
   if(p<=l and r<=q) {
      seg[node].lz = x;
      lazy(node, l, r);
      return;
   }
   int meio = (l+r)/2;
   update(2*node, l, meio, p, q, x);
   update(2*node+1, meio+1, r, p, q, x);
   seg[node].sum = seg[2*node].sum | seg[2*node+1].sum;
}

ll query(int node, int l, int r, int p, int q) {
   lazy(node, l, r);
   if(r<p or q<l) return 0;
   if(p<=l and r<=q) return seg[node].sum;
   int meio = (l+r)/2;
   return query(2*node, l, meio, p, q) | query(2*node+1, meio+1, r, p, q);
}
int main () {

   ios_base::sync_with_stdio(0);
   cin.tie(0);

   cin >> n >> q;

   for(int i = 1; i <= n; i++)
      cin >> v[i];

   for(int i = 1, a, b; i < n; i++) {
      cin >> a >> b;
      edge[a].pb(b);
      edge[b].pb(a);
   }

   dfs(1, 0);
   build(1, 1, n);

   while(q--) {
      int tp, node, color;
      cin >> tp;
      if(tp==1) {
         cin >> node >> color;
         update(1, 1, n, tin[node], tout[node], color);
      } else {
         cin >> node;
         ll bits = query(1, 1, n, tin[node], tout[node]);
         int resp = 0;
         for(ll i = ((ll)1<<60); i>=1; i/=2)
            if(bits>=i) bits-=i, resp++;
         cout << resp << "\n";
      }
   }
}

// 1 2 3 5 6 7 4